package com.xiao.tripmateai.advisor;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClientMessageAggregator;
import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
import reactor.core.publisher.Flux;

import java.time.Duration;
import java.time.Instant;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 自定义日志 Advisor
 * 打印 info 级别日志、只输出单次用户提示词和 AI 回复的文本
 */
@Slf4j
public class MyLoggerAdvisor implements CallAdvisor, StreamAdvisor {
    // 用于跟踪流式请求的上下文信息
    private final Map<String, Instant> streamStartTimes = new ConcurrentHashMap<>();

    @Override
    public String getName() {
        return this.getClass().getSimpleName();
    }

    @Override
    public int getOrder() {
        return 0;
    }

    /**
     * 请求与处理 ，在将对话发给ai之前做一些处理
     * @param request
     * @return
     */
    private ChatClientRequest before(ChatClientRequest request) {
        //log.info("AI Request，用户提示词: {}", request.prompt());

        //log.info("request.context().toString()" + request.context().toString());
        return request;
    }

    private void observeAfter(ChatClientResponse chatClientResponse) {
        String responseText = chatClientResponse.chatResponse().getResult().getOutput().getText();
        // 优化响应日志格式
        String cleanResponse = responseText
                .replace("\n", "\\n")  // 保留换行符标记
                .replaceAll("\\s{2,}", " "); // 压缩连续空格

        //log.info("AI回复: {}", cleanResponse);
    }

    @Override
    public ChatClientResponse adviseCall(ChatClientRequest chatClientRequest, CallAdvisorChain chain) {
        Instant start = Instant.now();
        chatClientRequest = before(chatClientRequest);
        ChatClientResponse chatClientResponse = chain.nextCall(chatClientRequest);
        observeAfter(chatClientResponse);
        // 记录同步调用总耗时
        Duration totalTime = Duration.between(start, Instant.now());
        log.info("同步调用总耗时: {}ms", totalTime.toMillis());
        return chatClientResponse;
    }

    @Override
    public Flux<ChatClientResponse> adviseStream(ChatClientRequest chatClientRequest, StreamAdvisorChain chain) {
        chatClientRequest = before(chatClientRequest);
        Flux<ChatClientResponse> chatClientResponseFlux = chain.nextStream(chatClientRequest);
        return (new ChatClientMessageAggregator()).aggregateChatClientResponse(chatClientResponseFlux, this::observeAfter);
    }
}
